"""
MHD模型推理示例
"""

import os
import sys
import argparse

# 添加项目根目录到路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from em_generation.config import get_config
from em_generation.inference import main as inference_main


def main():
    parser = argparse.ArgumentParser(description='MHD模型推理示例')
    parser.add_argument('--model_path', type=str, 
                        default='./checkpoints/mhd/best_model.pth', 
                        help='模型检查点路径')
    parser.add_argument('--data_path', type=str, 
                        default='./data/synthetic/mhd_data.h5', 
                        help='数据文件路径')
    parser.add_argument('--output_dir', type=str, default='./outputs/mhd/inference', 
                        help='输出目录')
    parser.add_argument('--device', type=str, default='auto', 
                        help='计算设备')
    parser.add_argument('--num_samples', type=int, default=5, 
                        help='推理样本数')
    parser.add_argument('--save_predictions', action='store_true', 
                        help='保存预测结果')
    parser.add_argument('--visualize', action='store_true', 
                        help='可视化结果')
    
    args = parser.parse_args()
    
    # 创建输出目录
    os.makedirs(args.output_dir, exist_ok=True)
    
    # 构建推理参数
    inference_args = [
        '--model_path', args.model_path,
        '--data_path', args.data_path,
        '--output_dir', args.output_dir,
        '--model_type', 'mhd',
        '--device', args.device,
        '--num_samples', str(args.num_samples)
    ]
    
    if args.save_predictions:
        inference_args.append('--save_predictions')
    
    if args.visualize:
        inference_args.append('--visualize')
    
    # 模拟命令行参数
    sys.argv = ['inference_mhd_example.py'] + inference_args
    
    print("开始MHD模型推理...")
    print(f"模型路径: {args.model_path}")
    print(f"数据路径: {args.data_path}")
    print(f"输出目录: {args.output_dir}")
    print(f"推理样本数: {args.num_samples}")
    print(f"保存预测结果: {args.save_predictions}")
    print(f"可视化结果: {args.visualize}")
    
    # 调用推理主函数
    inference_main()


if __name__ == "__main__":
    main()